import tensorflow.compat.v1 as tf
import tensorflow as tsf
from tensorflow.contrib.layers import fully_connected
from tensorflow.examples.tutorials.mnist import input_data
import numpy as np
from python_ai.common.xcommon import *

tf.set_random_seed(777)

'''
利用循环神经网络，完成手写体识别
'''
import sys
import os
path = r'../../../../large_data/DL1/mnist'
if not os.path.exists(path):
    print('[[[ DIR WRONG! ]]]', file=sys.stderr)
    sys.exit(0)
mnist = input_data.read_data_sets(path)
# mnist = input_data.read_data_sets(r"MNIST_data", one_hot=True)

check_shape(mnist.train.images, 'mnist.train.images')
check_shape(mnist.train.labels, 'mnist.train.labels')
check_shape(mnist.test.images, 'mnist.test.images')
check_shape(mnist.test.labels, 'mnist.test.labels')

# 参数
batch_size = 100 #mini-batch批大小
n_steps = 28 #时间步数（序列长度）
n_inputs = 28 #输入数据长度
n_neurons = 100 #隐藏状态，神经元个数
n_outputs = 10 #输出10分类
learning_rate = 0.001 #学习率
n_epochs = 1 #训练大周期
n_print = 20 # How many times in a epoch to print information

# 输入输出占位符
X = tf.placeholder(tf.float32, [None, n_steps, n_inputs]) #三维数据(?, 28, 28)
y = tf.placeholder(tf.int32, [None]) # 一维输出，标签0-9
# y = tf.placeholder(tf.float32, shape=[None, 10])

# 模型使用最简单的BasicRNNcell
basic_cell = tf.nn.rnn_cell.BasicRNNCell(num_units=n_neurons)
outputs, states = tf.nn.dynamic_rnn(basic_cell, X, dtype=tf.float32)  #outputs(?, 28, 100) states(?, 100)

logits = fully_connected(outputs[:, -1], n_outputs, activation_fn=None) #用最后一个时间步的输出
# w = tf.get_variable('w', [n_neurons, n_outputs], initializer=tsf.contrib.layers.xavier_initializer())
# b = tf.Variable(tf.random.normal([1, n_outputs]), dtype=tf.float32, name='b')
# logits = tf.matmul(states, w) + b

# 代价或损失函数
loss = tf.reduce_mean(tf.nn.sparse_softmax_cross_entropy_with_logits(logits=logits, labels=y))
# loss = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits_v2(logits=logits, labels=y))
training_op = tf.train.AdamOptimizer(learning_rate=learning_rate).minimize(loss)

# 计算准确率，只有等于y才是对的，其它都错
correct = tf.nn.in_top_k(logits, y, 1)  #每个样本的预测结果的前k个最大的数里面是否包含targets预测中的标签
accuracy = tf.reduce_mean(tf.cast(correct, tf.float32))
# acc_count = tf.equal(tf.argmax(logits, 1), tf.argmax(y, 1))
# accuracy = tf.reduce_mean(tf.cast(acc_count, dtype=tf.float32))

with tf.Session() as sess:
    sess.run(tf.global_variables_initializer())
    for epoch in range(n_epochs):
        n_batch = mnist.train.num_examples // batch_size
        group = n_batch // n_print
        for iteration in range(n_batch):  # 55000
            X_batch, y_batch = mnist.train.next_batch(batch_size)
            X_batch = X_batch.reshape((-1, n_steps, n_inputs)) #(?, 28, 28)
            _, loss_v, acc_v, outputs_v, states_v = sess.run([training_op, loss, accuracy, outputs, states], feed_dict={X: X_batch, y: y_batch})
            # print(np.unique(np.equal(outputs_v[:, -1], states_v)))
            if iteration % group == 0:
                print(f'#{epoch + 1}:{iteration + 1}: cost = {loss_v}, acc = {acc_v}')
        if iteration % group != 0:
            print(f'#{epoch + 1}:{iteration + 1}: cost = {loss_v}, acc = {acc_v}')
        acc_train = accuracy.eval(feed_dict={X: X_batch, y: y_batch})
        acc_test = accuracy.eval(feed_dict={X: mnist.test.images.reshape((-1, n_steps, n_inputs)),
                                            y: mnist.test.labels})
        print(epoch, "Train accuracy:", acc_train, "Test accuracy:", acc_test)
    print('Training over.')
